"""
Airflow环境中StarRocks数据库连接超时补丁
解决"Lost connection to MySQL server during query"问题

基于以下研究发现:
1. StarRocks虽然兼容MySQL协议，但在Airflow环境中需要特殊的超时设置
2. PyMySQL在没有明确超时的情况下可能会无限等待
3. Airflow的网络环境可能存在特殊的限制
"""

import logging
import pymysql
import socket

logger = logging.getLogger(__name__)

def patch_pymysql_for_starrocks():
    """
    为StarRocks数据库连接打补丁，解决Airflow环境中的超时问题
    """
    # 检查是否已经打过补丁
    if hasattr(pymysql.connect, '__starrocks_patched'):
        logger.info("🔧 StarRocks连接补丁已存在，跳过")
        return

    logger.info("🔧 开始应用StarRocks连接补丁...")
    
    # 保存原始连接函数
    original_connect = pymysql.connect
    
    def patched_connect(*args, **kwargs):
        """增强的PyMySQL连接函数，专门优化StarRocks连接"""
        
        # StarRocks优化连接参数
        starrocks_defaults = {
            'connect_timeout': 20,      # 连接超时
            'read_timeout': 60,         # 读取超时
            'write_timeout': 60,        # 写入超时
            'charset': 'utf8mb4',       # 字符集
            'autocommit': True,         # 自动提交
            'sql_mode': '',             # 清空SQL模式，避免兼容性问题
            'init_command': 'SET SESSION net_write_timeout=120',  # 增加写超时
        }
        
        # 合并参数（用户设置优先）
        for key, value in starrocks_defaults.items():
            if key not in kwargs:
                kwargs[key] = value
        
        # 特殊处理：如果没有设置数据库连接池参数，添加重试机制
        if 'database' in kwargs:
            logger.debug(f"🔗 连接StarRocks数据库: {kwargs.get('host', 'unknown')}:{kwargs.get('port', 'unknown')}/{kwargs['database']}")
        
        try:
            # 第一次尝试：使用标准参数
            return original_connect(*args, **kwargs)
            
        except (pymysql.err.OperationalError, socket.timeout, socket.error) as e:
            if "Lost connection" in str(e) or "timeout" in str(e).lower():
                logger.warning(f"⚠️ 首次连接失败，使用StarRocks优化参数重试: {str(e)}")
                
                # 第二次尝试：使用更保守的参数
                conservative_kwargs = kwargs.copy()
                conservative_kwargs.update({
                    'connect_timeout': 30,
                    'read_timeout': 120,
                    'write_timeout': 120,
                    'autocommit': False,  # 禁用自动提交
                })
                
                try:
                    return original_connect(*args, **conservative_kwargs)
                except Exception as e2:
                    logger.error(f"❌ StarRocks连接重试也失败: {str(e2)}")
                    raise e2
            else:
                raise e
    
    # 应用补丁
    pymysql.connect = patched_connect
    pymysql.connect.__starrocks_patched = True
    
    logger.info("✅ StarRocks连接补丁应用成功")

def patch_socket_defaults():
    """为socket设置默认超时，防止无限等待"""
    try:
        # 设置全局socket超时为60秒
        socket.setdefaulttimeout(60.0)
        logger.info("✅ 全局socket超时设置为60秒")
    except Exception as e:
        logger.warning(f"⚠️ 设置socket超时失败: {str(e)}")

def apply_all_patches():
    """应用所有连接补丁"""
    logger.info("🚀 开始应用Airflow-StarRocks连接补丁...")
    
    try:
        patch_pymysql_for_starrocks()
        patch_socket_defaults()
        logger.info("🎉 所有连接补丁应用完成")
        return True
    except Exception as e:
        logger.error(f"❌ 应用连接补丁失败: {str(e)}")
        return False

# 自动应用补丁（当模块被导入时）
if __name__ != "__main__":
    apply_all_patches()

if __name__ == "__main__":
    # 测试补丁
    print("🧪 测试StarRocks连接补丁...")
    apply_all_patches()
    print("✅ 补丁测试完成") 